{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "from forceFieldInfo import *\n",
    "# 从pdb文件读取原子类型和连接，构建id-atomtype list\n",
    "atomType = []\n",
    "connect = []\n",
    "\n",
    "with open('../../../sam2.pdb', 'r') as f:\n",
    "    for line in f:\n",
    "        if(line.startswith('HETATM')):\n",
    "            t = list(filter(lambda x: x, line.split(' ')))\n",
    "            atomType.append(t[2])\n",
    "        elif(line.startswith('CONECT')):\n",
    "            t = list(filter(lambda x: x, line.split(' ')))\n",
    "            connect.append(t[1:-1:])\n",
    "atomTypeStr = []\n",
    "for con in connect:\n",
    "    head = atomType[int(con[0])-1]\n",
    "    restList = [atomType[int(x)-1] for x in con[1::]]\n",
    "    restList.sort()\n",
    "    str1 = head + '_'+''.join(restList)\n",
    "    atomTypeStr.append(str1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "#读取坐标和 partial charge信息(h2o.chg)\n",
    "idx = 0\n",
    "atomInfo = []\n",
    "with open('../../../h2o.chg', 'r') as f:\n",
    "    for line in f.readlines():\n",
    "        t = list(filter(lambda x: x, line.split(' ')))\n",
    "        if t[0]!=atomType[idx]:\n",
    "            print('wrong atom type')\n",
    "        t[-1] =  '%f'%(float(t[-1])*1.2)\n",
    "        atomInfo.append(t)\n",
    "        idx += 1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "# bond\n",
    "bondList = []\n",
    "with open('../../../Data Bonds', 'r') as f:\n",
    "    for line in f.readlines():\n",
    "        t = list(filter(lambda x: x, line.replace('\\n','').split(' ')))\n",
    "        t2 = [atomTypeStr[int(x)-1] for x in t[2::]]\n",
    "        if t2[0]>t2[-1]:\n",
    "            t2.reverse()\n",
    "        t[1]=bondTypeDic[','.join(t2)]\n",
    "        bondList.append(t)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "# angle\n",
    "angleList = []\n",
    "with open('../../../Data Angles', 'r') as f:\n",
    "    for line in f.readlines():\n",
    "        t = list(filter(lambda x: x, line.replace('\\n', '').split(' ')))\n",
    "        t2 = [atomTypeStr[int(x)-1] for x in t[2::]]\n",
    "        if t2[0] > t2[-1]:\n",
    "            t2.reverse()\n",
    "        t[1] = angleTypeDic[','.join(t2)]\n",
    "        angleList.append(t)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dihedral\n",
    "dihedralList = []\n",
    "with open('../../../Data Dihedrals', 'r') as f:\n",
    "    for line in f.readlines():\n",
    "        t = list(filter(lambda x: x, line.replace('\\n', '').split(' ')))\n",
    "        t2 = [atomTypeStr[int(x)-1] for x in t[2::]]\n",
    "        if t2[0] > t2[-1]:\n",
    "            t2.reverse()\n",
    "        t3 = dihedralTypeDic.get(','.join(t2))\n",
    "        if t3 is not None:\n",
    "            t[1] = t3\n",
    "        else:\n",
    "            if t2[0].startswith('H') and t2[1].startswith('C') and t2[2].startswith('C') and t2[3].startswith('H'):\n",
    "                t[1] = '12'\n",
    "            elif t2[0].startswith('C') and t2[1].startswith('C') and t2[2].startswith('C') and t2[3].startswith('C') and len(t2[3]) == 6:\n",
    "                t[1] = '11'\n",
    "            elif t2[0].startswith('C') and t2[1].startswith('C') and t2[2].startswith('C') and t2[3].startswith('C') and len(t2[3]) == 5:\n",
    "                t[1] = '2'\n",
    "            elif t2[0].startswith('S') and t2[1].startswith('C') and t2[2].startswith('C') and t2[3].startswith('S'):\n",
    "                t[1] = '35'\n",
    "            elif t2[0].startswith('C') and t2[1].startswith('C') and t2[2].startswith('S') and t2[3].startswith('C'):\n",
    "                t[1] = '24'\n",
    "            elif t2[0].startswith('C') and t2[1].startswith('C') and t2[2].startswith('O') and t2[3].startswith('C'):\n",
    "                t[1] = '20'\n",
    "            else :\n",
    "                print(t2)\n",
    "        dihedralList.append(t)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "# improper\n",
    "improperList = []\n",
    "with open('../../../Data Impropers', 'r') as f:\n",
    "    for line in f.readlines():\n",
    "        t = list(filter(lambda x: x, line.replace('\\n', '').split(' ')))\n",
    "        t2 = [atomTypeStr[int(x)-1] for x in t[2::]]\n",
    "        if t2[0] > t2[-1]:\n",
    "            t2.reverse()\n",
    "        t3 = improperTypeDic.get(','.join(t2))\n",
    "        if t3 is not None:\n",
    "            t[1] = t3\n",
    "        else:\n",
    "            if t2[0].startswith('H') or t2[1].startswith('H') or t2[2].startswith('H') or t2[3].startswith('H'):\n",
    "                t[1] = '1'\n",
    "            else:\n",
    "                t[1] = '2'\n",
    "        improperList.append(t)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "samBoxAtomInfo = []\n",
    "with open('./sam_box.pdb', 'r') as f:\n",
    "    for line in f:\n",
    "        if(line.startswith('HETATM')):\n",
    "            t = list(filter(lambda x: x, line.split(' ')))\n",
    "            samBoxAtomInfo.append(t)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "#写lammps data 文件\n",
    "import copy\n",
    "numMolecular = 10\n",
    "numAl = 3072\n",
    "with open('./filler/system.data','w+') as f:\n",
    "    f.write('LAMMPS Description\\n\\n')\n",
    "    #原子，键，键角，二面角，共面项数量\n",
    "    f.write('    %d  atoms\\n' % (len(atomType)*numMolecular+numAl))\n",
    "    f.write('    %d  bonds\\n' % (len(bondList)*numMolecular))\n",
    "    f.write('    %d  angles\\n' % (len(angleList)*numMolecular))\n",
    "    f.write('    %d  dihedrals\\n' % (len(dihedralList)*numMolecular))\n",
    "    f.write('    %d  impropers\\n\\n' % (len(improperList)*numMolecular))\n",
    "\n",
    "    #原子，键，键角，二面角，共面项 类型数量\n",
    "    f.write('    %d  atom types\\n' % (len(atomType)+2))\n",
    "    f.write('    %d  bond types\\n' % (len(bondCoeffDic)))\n",
    "    f.write('    %d  angle types\\n' % (len(angleCoeffDic)))\n",
    "    f.write('    %d  dihedral types\\n' % (len(dihedralCoeffDic)))\n",
    "    f.write('    %d  improper types\\n\\n' % (len(improperCoeffDic)))\n",
    "\n",
    "    #箱盒大小\n",
    "    f.write('    0.0 32.382400 xlo xhi\\n')\n",
    "    f.write('    0.0 32.382400 ylo yhi\\n')\n",
    "    f.write('    -51 53 zlo zhi\\n')\n",
    "\n",
    "    f.write('\\n')\n",
    "\n",
    "    #质量\n",
    "    f.write('Masses\\n\\n')\n",
    "    for i in range(len(atomType)):\n",
    "        if(atomType[i]=='C'):\n",
    "            m = '12.01'\n",
    "        elif(atomType[i] == 'H'):\n",
    "            m = '1.008'\n",
    "        elif(atomType[i] == 'H'):\n",
    "            m = '14.01'\n",
    "        elif(atomType[i] == 'O'):\n",
    "            m = '16.00'\n",
    "        elif(atomType[i] == 'S'):\n",
    "            m = '32.06'\n",
    "        f.write('%d %s  # %s\\n'%(i+1,m,atomTypeStr[i]))\n",
    "    f.write('%d 26.982000\\n' %(len(atomType)+1))\n",
    "    f.write('%d 26.982000\\n' % (len(atomType)+2))\n",
    "    f.write('\\n')\n",
    "\n",
    "    #原子参数,full类型\n",
    "    #atom-ID molecule-ID atom-type q(partial charge) x y z\n",
    "    f.write('Atoms  # full\\n\\n')\n",
    "    for j in range(numMolecular):\n",
    "        for i in range(len(atomType)):\n",
    "            f.write('%d %d %d %s %s %s %s\\n' %\n",
    "                    (j*444+i+1, j+1, i+1, atomInfo[i][4], samBoxAtomInfo[j*444+i][5], samBoxAtomInfo[j*444+i][6], samBoxAtomInfo[j*444+i][7]))\n",
    "    AlData = []\n",
    "    numSamAtom = len(atomType)*numMolecular\n",
    "    with open('./al.base.lmp','r') as fal:\n",
    "        for line in fal.readlines():\n",
    "            t = line.split('\\t')\n",
    "            t[0] = str(int(t[0])+numSamAtom)\n",
    "            z = float(t[-1])\n",
    "            if z <0.1 or z > 46:\n",
    "                alAtomType = 445\n",
    "            else:\n",
    "                alAtomType = 446\n",
    "            f.write('%s %d %d %s %s %s %f\\n' %\n",
    "                    (t[0], 11, alAtomType, t[3], t[4], t[5], z-51))\n",
    "    f.write('\\n')\n",
    "\n",
    "    #bonds\n",
    "    f.write('Bonds\\n\\n')\n",
    "    for j in range(numMolecular):\n",
    "        for t in bondList:\n",
    "            tt = copy.deepcopy(t)\n",
    "            tt[0] = int(tt[0])+len(bondList)*j\n",
    "            tt[2] = int(tt[2])+444*j\n",
    "            tt[3]=int(tt[3])+444*j\n",
    "            f.write(' '.join(map(str,tt)))\n",
    "            f.write('\\n')\n",
    "    f.write('\\n')\n",
    "\n",
    "    #angles\n",
    "    f.write('Angles\\n\\n')\n",
    "    for j in range(numMolecular):\n",
    "        for t in angleList:\n",
    "            tt = copy.deepcopy(t)\n",
    "            tt[0] = int(tt[0])+len(angleList)*j\n",
    "            tt[2] = int(tt[2])+444*j\n",
    "            tt[3] = int(tt[3])+444*j\n",
    "            tt[4]=int(tt[4])+444*j\n",
    "            f.write(' '.join(map(str, tt)))\n",
    "            f.write('\\n')\n",
    "    f.write('\\n')\n",
    "\n",
    "    #dihedral\n",
    "    f.write('Dihedrals\\n\\n')\n",
    "    for j in range(numMolecular):\n",
    "        for t in dihedralList:  \n",
    "            tt = copy.deepcopy(t)\n",
    "            tt[0]=int(tt[0])+len(dihedralList)*j\n",
    "            tt[2] = int(tt[2])+444*j\n",
    "            tt[3] = int(tt[3])+444*j\n",
    "            tt[4] = int(tt[4])+444*j\n",
    "            tt[5]=int(tt[5])+444*j\n",
    "            f.write(' '.join(map(str, tt)))\n",
    "            f.write('\\n')\n",
    "    f.write('\\n')\n",
    "\n",
    "    #improper\n",
    "    f.write('Impropers\\n\\n')\n",
    "    for j in range(numMolecular):\n",
    "        for t in improperList:\n",
    "            tt = copy.deepcopy(t)\n",
    "            tt[0]=int(tt[0])+len(improperList)*j\n",
    "            tt[2] = int(tt[2])+444*j\n",
    "            tt[3] = int(tt[3])+444*j\n",
    "            tt[4] = int(tt[4])+444*j\n",
    "            tt[5] = int(tt[5])+444*j\n",
    "            f.write(' '.join(map(str, tt)))\n",
    "            f.write('\\n')\n",
    "    f.write('\\n')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 写LAMMPS力场参数文件\n",
    "with open('./filler/system.in.coeff', 'w+') as f:\n",
    "    #写pair coeff\n",
    "    for i in range(len(atomTypeStr)):\n",
    "        coeff = nonBondCoeffDic[atomTypeStr[i]].split(',')\n",
    "        f.write('    pair_coeff  %d %d lj/cut/coul/long %s %s  #%s\\n' %\n",
    "                (i+1, i+1, coeff[0], coeff[1], atomTypeStr[i]))\n",
    "        ljCutCoeff0 = (float(coeff[0])*0.31)**0.5\n",
    "        ljCutCoeff1 = (float(coeff[1])*3.893227)**0.5\n",
    "        f.write('    pair_coeff %d %s lj/cut %f %f  #%s-Al\\n' %\n",
    "                (i+1, '445*446', ljCutCoeff0, ljCutCoeff1, atomTypeStr[i]))\n",
    "    f.write('    pair_coeff 445*446 445*446 eam Al_jnp.eam # al-al\\n')\n",
    "    #f.write('    pair_coeff 1*444 446 none # al-al\\n')\n",
    "    #bond coeff\n",
    "    for key in bondCoeffDic.keys():\n",
    "        f.write('    bond_coeff %s  %s\\n' %\n",
    "                (key, bondCoeffDic[key].replace(',', ' ')))\n",
    "    #angle coeff\n",
    "    for key in angleCoeffDic.keys():\n",
    "        f.write('    angle_coeff %s  %s\\n' %\n",
    "                (key, angleCoeffDic[key].replace(',', ' ')))\n",
    "    # dihedral coeff\n",
    "    for key in dihedralCoeffDic.keys():\n",
    "        f.write('    dihedral_coeff %s  %s\\n' %\n",
    "                (key, dihedralCoeffDic[key].replace(',', ' ')))\n",
    "    # improper coeff\n",
    "    for key in improperCoeffDic.keys():\n",
    "        f.write('    improper_coeff %s  %s\\n' %\n",
    "                (key, improperCoeffDic[key].replace(',', ' ')))\n"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "657c0b48f02cef986c54fd63ced157a09a9f13d09ea3bbd6d5802a6509aaddf2"
  },
  "kernelspec": {
   "display_name": "Python 3.8.3 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
